#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2021 wanghch <wanghch@wanghch-pc>
#
# Distributed under terms of the MIT license.

"""

"""
from jqdatasdk import *
from jq_utils import *
import pandas as pd
import datetime
import os
import sys
today = datetime.datetime.today()

sdk_auth()

ctype = "st"
if len(sys.argv) == 2:
    ctype = sys.argv[1]


all_days = get_trade_days("2001-01-01", today.strftime("%Y-%m-%d"))
#print(all_days)

output_dir = "data"
if ctype == "st":
    df = pd.read_csv("./clist.csv")
elif ctype == "index":
    output_dir = "dataindex"
    df = pd.read_csv("./index.csv")

code_map = {}
for index, row in df.iterrows():
    code_map[row['code']] = row['display_name']

codes = list(code_map.keys())

fields = ['open','close','low','high','volume','money']
csv_fields = ['date'] + fields

last_day = all_days[-1].strftime("%Y-%m-%d")
dfs = []
for code, name in code_map.items():
    dfile = output_dir + "/" + code + "_" + name +".csv"
    start_date = all_days[0]
    incr_flag = False
    dfd = None
    if os.path.exists(dfile):
        dfd = pd.read_csv(dfile)
        dc = dfd.columns[0]
        if dc != "date":
            dfd.rename(columns = {dc: "date"},  inplace=True)


        # print(dfd['date'].max())
        #dfd['date'] = pd.to_datetime(dfd['date'])
        dmax = dfd.date.max()
        # print(dfile, dmax)
        #sdmax = dmax.strftime("%Y-%m-%d")
        # print(type(all_days[-1]))
        if last_day != dmax:
            start_date = datetime.datetime.strptime(dmax, '%Y-%m-%d') + datetime.timedelta(days = 1)
            # print(start_date, all_days[-1])
            incr_flag = True
        else:
            dfd[dfd.close > 0].to_csv(dfile, index = False)
            continue
    df = get_price(code, start_date = start_date,  end_date=all_days[-1], fq='pre', frequency='1d', fields=fields)
    df['name'] = name
    df['code'] = code
    if incr_flag and (dfd is not None):
        print("incr ", start_date, " ", code, " ", name)
        df['date'] = df.index
        df.index = range(df.shape[0])
        df['date'] = df['date'].apply(lambda x: x.strftime("%Y-%m-%d"))
        pd.concat([dfd, df], ignore_index = True, sort=False).to_csv(dfile, index = False)
    else:
        df[df.close > 0].to_csv(dfile)

